/**
 * @file main.cpp
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "softmaxv2_tiling.h"
#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
#include "aclrtlaunch_softmaxv2_custom.h"
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void softmaxv2_custom(GM_ADDR srcGM, GM_ADDR dstGM, GM_ADDR softmaxMaxGM, 
                                                        GM_ADDR softmaxSumGM, GM_ADDR tilingGM);
#endif

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 1;
    size_t oriRSize = 4096+ 5;
    size_t oriASize = 4;
    size_t perElementByBlk = 8;
    size_t srcFileSize = oriRSize * oriASize * sizeof(float);
    size_t softmaxMaxFileSize = oriASize * perElementByBlk * sizeof(float);
    size_t softmaxTilingFileSize = sizeof(SoftmaxV2TilingData);

#ifdef ASCENDC_CPU_DEBUG
    uint8_t *tiling = (uint8_t *)AscendC::GmAlloc(tilingSize);
    ReadFile("./input/input_tiling.bin", tilingSize, tiling, tilingSize);
    uint8_t *x = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *y = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *z = (uint8_t *)AscendC::GmAlloc(outputByteSize);

    ReadFile("./input/input_x.bin", inputByteSize, x, inputByteSize);
    ReadFile("./input/input_y.bin", inputByteSize, y, inputByteSize);

    AscendC::SetKernelMode(KernelMode::AIV_MODE);

    ICPU_RUN_KF(add_custom, blockDim, x, y, z,
                *reinterpret_cast<AddCustomTilingData *>(tiling)); // use this macro for cpu debug

    WriteFile("./output/output_z.bin", z, outputByteSize);

    AscendC::GmFree((void *)x);
    AscendC::GmFree((void *)y);
    AscendC::GmFree((void *)z);
    AscendC::GmFree((void *)tiling);
#else
    // const char *aclConfigPath = "./acl.json";
    // CHECK_ACL(aclInit(aclConfigPath));
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    uint8_t *srcHost;
    uint8_t *srcDevice;
    CHECK_ACL(aclrtMallocHost((void **)(&srcHost), srcFileSize));
    CHECK_ACL(aclrtMalloc((void **)&srcDevice, srcFileSize,
                        ACL_MEM_MALLOC_HUGE_FIRST));
    ReadFile("./input/src.bin", srcFileSize, srcHost, srcFileSize);
    CHECK_ACL(aclrtMemcpy(srcDevice, srcFileSize, srcHost, srcFileSize,
                        ACL_MEMCPY_HOST_TO_DEVICE));

    uint8_t *dstHost;
    uint8_t *dstDevice;
    CHECK_ACL(aclrtMallocHost((void **)(&dstHost), srcFileSize));
    CHECK_ACL(aclrtMalloc((void **)&dstDevice, srcFileSize,
                        ACL_MEM_MALLOC_HUGE_FIRST));   

    uint8_t *softmaxMaxHost;
    uint8_t *softmaxMaxDevice;
    CHECK_ACL(aclrtMallocHost((void **)(&softmaxMaxHost), softmaxMaxFileSize));
    CHECK_ACL(aclrtMalloc((void **)&softmaxMaxDevice, softmaxMaxFileSize,
                        ACL_MEM_MALLOC_HUGE_FIRST));   
    
    uint8_t *softmaxSumHost;
    uint8_t *softmaxSumDevice;
    CHECK_ACL(aclrtMallocHost((void **)(&softmaxSumHost), softmaxMaxFileSize));
    CHECK_ACL(aclrtMalloc((void **)&softmaxSumDevice, softmaxMaxFileSize,
                        ACL_MEM_MALLOC_HUGE_FIRST));   

    uint8_t *softmaxTilingHost;
    uint8_t *softmaxTilingDevice;
    CHECK_ACL(aclrtMallocHost((void **)(&softmaxTilingHost), softmaxTilingFileSize));
    CHECK_ACL(aclrtMalloc((void **)&softmaxTilingDevice, softmaxTilingFileSize, ACL_MEM_MALLOC_HUGE_FIRST));
    ReadFile("./input/tiling.bin", softmaxTilingFileSize, softmaxTilingHost, softmaxTilingFileSize);
    CHECK_ACL(aclrtMemcpy(softmaxTilingDevice, softmaxTilingFileSize, softmaxTilingHost, softmaxTilingFileSize, ACL_MEMCPY_HOST_TO_DEVICE));

    ACLRT_LAUNCH_KERNEL(softmaxv2_custom)(blockDim, stream, srcDevice, dstDevice, softmaxMaxDevice, softmaxSumDevice, softmaxTilingDevice);
    aclGetRecentErrMsg();
    std::cout << aclGetRecentErrMsg() << std::endl;
    CHECK_ACL(aclrtSynchronizeStream(stream));
    std::cout << aclGetRecentErrMsg() << std::endl;
    CHECK_ACL(aclrtMemcpy(softmaxMaxHost, softmaxMaxFileSize, softmaxMaxDevice, softmaxMaxFileSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/gen_softmaxMax.bin", softmaxMaxHost, softmaxMaxFileSize);
    CHECK_ACL(aclrtMemcpy(softmaxSumHost, softmaxMaxFileSize, softmaxSumDevice, softmaxMaxFileSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/gen_softmaxSum.bin", softmaxSumHost, softmaxMaxFileSize);
    CHECK_ACL(aclrtMemcpy(dstHost, srcFileSize, dstDevice, srcFileSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/gen_dst.bin", dstHost, srcFileSize);

    CHECK_ACL(aclrtFree(srcDevice));
    CHECK_ACL(aclrtFree(dstDevice));
    CHECK_ACL(aclrtFree(softmaxMaxDevice));
    CHECK_ACL(aclrtFree(softmaxSumDevice));
    CHECK_ACL(aclrtFree(softmaxTilingDevice));

    CHECK_ACL(aclrtFreeHost(srcHost));
    CHECK_ACL(aclrtFreeHost(dstHost));
    CHECK_ACL(aclrtFreeHost(softmaxMaxHost));
    CHECK_ACL(aclrtFreeHost(softmaxSumHost));
    CHECK_ACL(aclrtFreeHost(softmaxTilingHost));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}
